package com.carp.shard.plugin.strategy;

import cn.hutool.core.util.ReflectUtil;
import com.alibaba.druid.sql.ast.statement.SQLExprTableSource;
import com.alibaba.druid.sql.visitor.SchemaStatVisitor;
import com.alibaba.druid.stat.TableStat;
import com.alibaba.fastjson.JSON;
import com.carp.shard.plugin.config.MyBatisShardProperties;
import com.carp.shard.plugin.constant.ReadType;
import com.carp.shard.plugin.constant.WriteType;
import com.carp.shard.plugin.exception.MyBatisShardException;
import lombok.Getter;
import lombok.NonNull;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.reflection.MetaObject;
import org.springframework.beans.factory.annotation.Autowired;

import javax.annotation.PostConstruct;
import java.util.*;
import java.util.function.Supplier;

/**
 * 抽象分片策略
 *
 * @author lile
 * @since 2024/5/10
 */
@Getter
public abstract class AbstractShardStrategy {

    private ReadType readType = ReadType.OLD;
    private WriteType writeType = WriteType.OLD;

    @Autowired
    private MyBatisShardProperties myBatisShardProperties;

    /**
     * 获取原始表名
     *
     * @return 原始表名
     */
    protected abstract String getOriginalTableName();

    /**
     * 获取分片表名
     *
     * @param parameterObject 参数
     * @return 分片表名
     */
    protected abstract String getShardTableName(@NonNull Object parameterObject);

    /**
     * 替换表名
     *
     * @param x        SQLExprTableSource AST节点
     * @param boundSql boundSql
     * @return 是否有替换
     */
    public abstract boolean replaceTableName(@NonNull SQLExprTableSource x, @NonNull BoundSql boundSql);

    /**
     * 处理参数
     *
     * @param metaObject        metaObject
     * @param boundSql          boundSql
     * @param schemaStatVisitor 用于统计SQL中使用的表、字段、过滤条件、排序表达式、分组表达式等的Visitor
     */
    public abstract void processParams(@NonNull MetaObject metaObject, @NonNull BoundSql boundSql, @NonNull SchemaStatVisitor schemaStatVisitor);

    public AbstractShardStrategy() {
        register();
    }


    /**
     * 加载分片配置
     */
    @PostConstruct
    private void loadConfig() {
        MyBatisShardProperties.ShardProperty shardProperty = myBatisShardProperties
                .getShardPropertyMap().get(getOriginalTableName());
        if (Objects.isNull(shardProperty)) {
            return;
        }
        readType = Optional.ofNullable(shardProperty.getReadType()).orElse(ReadType.OLD);
        writeType = Optional.ofNullable(shardProperty.getWriteType()).orElse(WriteType.OLD);
    }

    /**
     * 注册到分片策略上下文
     */
    private void register() {
        ShardStrategyContext.registerStrategy(getOriginalTableName(), this);
    }

    /**
     * 给指定列赋值
     *
     * @param boundSql   boundSql
     * @param columnName 列名
     * @param columns    SQL中出现的列
     * @param supplier   值的supplier
     */
    @SuppressWarnings({"rawtypes", "unchecked"})
    protected void setColumnValue(BoundSql boundSql, String columnName,
                                  Collection<TableStat.Column> columns, Supplier<Object> supplier) {

        boolean findColumn = false;
        for (TableStat.Column column : columns) {
            if (Objects.equals(column.getName(), columnName)) {
                findColumn = true;
                break;
            }
        }
        if (!findColumn) {
            return;
        }

        Object parameterObject = boundSql.getParameterObject();
        if (parameterObject instanceof Map) {
            if (((Map) parameterObject).containsKey(columnName)) {
                ((Map) parameterObject).computeIfAbsent(columnName, k -> supplier.get());
                return;
            }
            Set<Map.Entry> entrySet = ((Map) parameterObject).entrySet();
            for (Map.Entry entry : entrySet) {
                Object value = entry.getValue();
                if (value instanceof Collection) {
                    for (Object o : (Collection) value) {
                        setNx(o, columnName, supplier);
                    }
                    return;
                }
            }
            return;
        }
        setNx(parameterObject, columnName, supplier);
    }

    private void setNx(@NonNull Object o, @NonNull String key, Supplier<Object> supplier) {
        if (Objects.nonNull(ReflectUtil.getFieldValue(o, key))) {
            return;
        }
        ReflectUtil.setFieldValue(o, key, supplier.get());
    }

    @SuppressWarnings({"rawtypes", "unchecked"})
    protected Object getValue(@NonNull Object parameterObject, @NonNull String key) {
        if (parameterObject instanceof Map) {
            if (((Map) parameterObject).containsKey(key)) {
                return ((Map) parameterObject).get(key);
            }
            Set<Map.Entry> entrySet = ((Map) parameterObject).entrySet();
            for (Map.Entry entry : entrySet) {
                Object value = entry.getValue();
                if (ReflectUtil.hasField(value.getClass(), key)) {
                    return ReflectUtil.getFieldValue(value, key);
                }
                if (value instanceof Collection) {
                    for (Object o : (Collection) value) {
                        return ReflectUtil.getFieldValue(o, key);
                    }
                }
            }
            throwEx(parameterObject, key);
        }

        return ReflectUtil.getFieldValue(parameterObject, key);
    }

    private void throwEx(@NonNull Object parameterObject, @NonNull String key) {
        throw new MyBatisShardException(String.format("Error: Method AbstractShardStrategy.getValue, parameterObject(%s): %s, key: %s", parameterObject.getClass().getSimpleName(), JSON.toJSONString(parameterObject), key));
    }

}
